from array import *
import numpy as np
import csv
import matplotlib.pyplot as plt
import math as m
import pandas as pd

from scipy.fftpack import fft, ifft



def read_data(f):
    print "Reading file...",
    fo1=open(f, 'r')
    fo =csv.reader(fo1,delimiter=' ')

    hns=array('f')
    hew=array('f')
    v=array('f')
    a=0
    for i in fo:
        hns.append(float(i[3]))
        hew.append(float(i[5]))
        v.append(float(i[7]))
        a=a+1
    fo1.close()
    print "finished."
    return np.asarray(hns),np.asarray(hew),np.asarray(v)



def plot_TS(ts):
    #plotting time series 
    print "Plotting..."
    fig=plt.figure()
    ax=fig.add_subplot(111)
    for g in range(len(ts.columns)):
        ax.plot(ts.index,(ts[ts.columns[g]].values/np.nanmax(ts[ts.columns[g]].values))+(-1*g), label=ts.columns[g])
        
    ax.set_yticklabels([],visible=False)
    ax.set_xlabel('Time (s)')
    ax.legend(loc='upper left')
    ax.set_title('Normalized component time series')
    plt.grid(b=None, which='both', axis='both')
    return





def compute_plot_HVSR_depth(fd,freq,winsize,Vs_list):
    #smoothing using moving average
    if winsize<=0 or winsize>20:
        fd_av=fd
    else:
        print "Smoothing data..."
        fd_av=pd.rolling_mean(fd,int(len(freq)*(winsize/100.)))

    #computing HVSR
    print "Computing HVSR..."
    hvsr=np.sqrt(fd_av['NS'].values**2+fd_av['NS'].values**2)/fd_av['V'].values
    fd_av['HVSR']=hvsr

    #spectral plot (frequency)
    print "Plotting..."
    fig=plt.figure()
    ax=fig.add_subplot(111)
    for g in range(len(fd_av.columns)):
        ax.plot(fd_av.index,(fd_av[fd_av.columns[g]].values/np.nanmax(fd_av[fd_av.columns[g]].values))+(-1*g), label=fd_av.columns[g])
        
    ax.set_yticklabels([],visible=False)
    ax.set_xlabel('Frequency (Hz)')
    ax.semilogx()
    ax.legend(loc='upper left')
    ax.set_title('Normalized component frequency spectra and HVSR')
    plt.grid(b=None, which='both', axis='both')

    #spectral plot (depth)
    fig=plt.figure()
    ax=fig.add_subplot(111)

    for v in range(len(Vs_list)):
        H=Vs_list[v]/(4.*fd_av.index.values)
        hlab="Vs="+str(Vs_list[v])
        ax.plot(H,(hvsr/np.nanmax(hvsr))+(-1*v), label=hlab)

    ax.set_yticklabels([],visible=False)
    ax.set_xlabel('Equivalent depth (m)')
    ax.semilogx()
    ax.legend(loc='upper right')
    ax.set_title('Normalized HVSR for various Vs (m/s)')
    plt.grid(b=None, which='both', axis='both')

    

    return

    



#MAIN

#INPUTS
samplingrate=100.  #number of samples per second
Vs_list=(100,200,300,400,500)  #m/s
inputfilename='ASCII_Waveform_Data.txt'  #textfile from OYO microtremor instrument processed in PickWin

#READ INPUT FILE
hns,hew,v=read_data(inputfilename)

#creating dataframe of time domain data
data={'NS':hns,'EW':hew,'V':v}
ts=pd.DataFrame(data=data,index=np.arange(len(hns))/samplingrate)

#COMPUTING FFT
freq=np.arange(len(hns)/2)/(len(hns)/samplingrate)
hns_fft=np.abs(fft(hns)[1:len(freq)])
hew_fft=np.abs(fft(hew)[1:len(freq)])
v_fft=np.abs(fft(v)[1:len(freq)])

#creating dataframe of frequency domain data
data={'NS':hns_fft,'EW':hew_fft,'V':v_fft}
fd=pd.DataFrame(data=data,index=freq[1:])

change_smoothing_window=1
while change_smoothing_window==1:
    plot_TS(ts)
    winsize=float(raw_input("Input size of smoothing window, as percent of data length (0-20): "))
    print ""
    compute_plot_HVSR_depth(fd,freq,winsize,Vs_list)
    plt.show()
    print ""
    change_smoothing_window=int(raw_input("Replot using another smoothing window?: 0 <NO>     1 <YES>"))


                     

